import os
import torch
import shutil
import fastapi
import typing
import tempfile
from service.config import config
from scripts import crop_and_ocr


route_model = fastapi.APIRouter(tags=["model"])


@route_model.post('/upload-model')
def upload_model(model_file: fastapi.UploadFile = fastapi.File(...)):
    global model
    try:
        model_path = os.path.join(tempfile.gettempdir(), model_file.filename)
        with open(model_path, "wb") as buffer:
            shutil.copyfileobj(model_file.file, buffer)
        model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
        return {
            "message": "Model uploaded successfully",
            "model_path": model_path}
    except Exception as e:
        raise fastapi.HTTPException(status_code=500, detail=str(e))


@route_model.post('/process-image')
def process_image(file: fastapi.UploadFile = fastapi.File(...),
                  ocr_model_flag: typing.Optional[str] = fastapi.Form(
        default=config()['ORC_MODEL_FLAG']),
        confidence_threshold: typing.Optional[float] = fastapi.Form(
    default=config()['CONFIDENCE_THRESHOLD']),
        iou_threshold: typing.Optional[float] = fastapi.Form(
    default=config()['IOU_THRESHOLD']
)):
    global model, processed_text_data
    if model is None:
        raise fastapi.HTTPException(
            status_code=400, detail="Model not loaded. Please upload the model first.")

    try:
        image_path = os.path.join(tempfile.gettempdir(), file.filename)

        with open(image_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)

        processed_text_data = crop_and_ocr(
            image_path, model, confidence_threshold, iou_threshold, ocr_model_flag)

        return {"text_results": processed_text_data}
    except Exception as e:
        raise fastapi.HTTPException(
            status_code=500, detail=str(e))
